package com.ai.ai.controller;

import jakarta.servlet.http.HttpServletResponse;
import org.apache.commons.lang3.StringUtils;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.Iterator;
import java.util.Map;

@RequestMapping("/chat")
@RestController
public class ChatController {

    private final ChatModel chatModel;

    public ChatController(OpenAiChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @GetMapping("/call")
    public Map<String, String> call(@RequestParam String message) {
        return Map.of("generation", this.chatModel.call(message));
    }

    @GetMapping(value = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> stream(@RequestParam String message, HttpServletResponse response) {
        response.setCharacterEncoding("UTF-8");
        Prompt prompt = new Prompt(new UserMessage(message));
        Flux<ChatResponse> flux = this.chatModel.stream(prompt);
        return flux.map(chatResponse -> {
            StringBuffer result = new StringBuffer();
            if (chatResponse != null && chatResponse.getResults() != null && chatResponse.getResults().size() > 0) {
                Iterator<Generation> iterator = chatResponse.getResults().iterator();
                while (iterator.hasNext()) {
                    Generation generation = iterator.next();
                    if (generation.getOutput() != null && generation.getOutput().getText() != null) {
                        result.append(generation.getOutput().getText());
                    }
                }
            }
            return result.toString();
        }
        ).filter(StringUtils::isNotEmpty).doOnComplete(new Runnable() {
            @Override
            public void run() {
                try {
                    response.getOutputStream().close();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
        }).onErrorComplete(Throwable.class);
    }
}
